package security

import (
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"github.com/emmansun/gmsm/padding"
	"github.com/emmansun/gmsm/sm4"
	"io"
)

// SM4KEY SM4 key 32 bytes
var SM4KEY, _ = hex.DecodeString("1234567890abcdef1234567890abcdef")

func EncryptSM4(data string) (string, error) {
	block, err := sm4.NewCipher(SM4KEY)
	if err != nil {
		return "", err
	}
	pkcs7 := padding.NewPKCS7Padding(sm4.BlockSize)
	paddedPlainText := pkcs7.Pad([]byte(data))
	ciphertext := make([]byte, sm4.BlockSize+len(paddedPlainText))
	iv := ciphertext[:sm4.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return "", err
	}
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext[sm4.BlockSize:], paddedPlainText)
	return hex.EncodeToString(ciphertext), nil
}

func DecryptSM4(data string) (string, error) {
	block, err := sm4.NewCipher(SM4KEY)
	if err != nil {
		return "", err
	}
	ciphertext, err := hex.DecodeString(data)
	if err != nil {
		return "", err
	}
	if len(ciphertext) < sm4.BlockSize {
		return "", err
	}
	iv := ciphertext[:sm4.BlockSize]
	ciphertext = ciphertext[sm4.BlockSize:]
	if len(ciphertext)%sm4.BlockSize != 0 {
		return "", err
	}
	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(ciphertext, ciphertext)
	pkcs7 := padding.NewPKCS7Padding(sm4.BlockSize)
	plainText, err := pkcs7.Unpad(ciphertext)
	if err != nil {
		return "", err
	}
	return string(plainText), nil
}
